import matplotlib.pyplot as plt
import numpy as np
import logging
from scipy.signal import savgol_filter as sa
import sys


class Plot:
    def __init__(self,
                 model: str,
                 dataset: str):
        self.model = model
        self.dataset = dataset
        flag = False
        try:
            self.baseline_loss = np.load(f'./res/data/{self.model}_{self.dataset}_None_loss.npy')
            self.baseline_acc1 = np.load(f'./res/data/{self.model}_{self.dataset}_None_acc1.npy')
            self.baseline_acc5 = np.load(f'./res/data/{self.model}_{self.dataset}_None_acc5.npy')
        except:
            logging.error(f"Please conduct {self.model} on {self.dataset} without compression")
            flag = True

        try:
            self.quan2_loss = np.load(f'./res/data/{self.model}_{self.dataset}_2-bit quantization_loss.npy')
            self.quan2_acc1 = np.load(f'./res/data/{self.model}_{self.dataset}_2-bit quantization_acc1.npy')
            self.quan2_acc5 = np.load(f'./res/data/{self.model}_{self.dataset}_2-bit quantization_acc5.npy')
        except:
            logging.error(f"Please conduct {self.model} on {self.dataset} with 2-bit quantization")
            flag = True

        try:
            self.quan8_loss = np.load(f'./res/data/{self.model}_{self.dataset}_8-bit quantization_loss.npy')
            self.quan8_acc1 = np.load(f'./res/data/{self.model}_{self.dataset}_8-bit quantization_acc1.npy')
            self.quan8_acc5 = np.load(f'./res/data/{self.model}_{self.dataset}_8-bit quantization_acc5.npy')
        except:
            logging.error(f"Please conduct {self.model} on {self.dataset} with 8-bit quantization")
            flag = True

        try:
            self.sparse_loss = np.load(f'./res/data/{self.model}_{self.dataset}_Top-2% sparsification_loss.npy')
            self.sparse_acc1 = np.load(f'./res/data/{self.model}_{self.dataset}_Top-2% sparsification_acc1.npy')
            self.sparse_acc5 = np.load(f'./res/data/{self.model}_{self.dataset}_Top-2% sparsification_acc5.npy')
        except:
            logging.error(f"Please conduct {self.model} on {self.dataset} with Top-2% sparsification")
            flag = True

        try:
            t = f"Semi-sparsified (Top-2%) Features with Dynamic Mask Encoding (2-bit)"
            self.ssfdme_loss = np.load(f'./res/data/{self.model}_{self.dataset}_{t}_loss.npy')
            self.ssfdme_acc1 = np.load(f'./res/data/{self.model}_{self.dataset}_{t}_acc1.npy')
            self.ssfdme_acc5 = np.load(f'./res/data/{self.model}_{self.dataset}_{t}_acc5.npy')
        except:
            logging.error(f"Please conduct {self.model} on {self.dataset} with "
                          f"Semi-sparsified (Top-2%) Features with Dynamic Mask Encoding (2-bit)")
            flag = True

        if flag:
            sys.exit(0)

    def plot(self, fontsize=12, step=100):
        self.baseline_acc1 *= 100
        self.quan2_acc1 *= 100
        self.quan8_acc1 *= 100
        self.sparse_acc1 *= 100
        self.ssfdme_acc1 *= 100
        fig, ax = plt.subplots(1, 2, figsize=(10, 3), dpi=600)
        ax[0].plot(sa(self.baseline_acc1, 3, 1), "-", label="Baseline")
        ax[0].plot(sa(self.quan2_acc1, 3, 1), "--", label="2-bit quantization")
        ax[0].plot(sa(self.quan8_acc1, 3, 1), ":", label="8-bit quantization")
        ax[0].plot(sa(self.sparse_acc1, 3, 1), "-.", label="98% sparsification")
        ax[0].plot(sa(self.ssfdme_acc1, 3, 1), ".", label="Ours")
        # ax[0].legend(prop={'size':6})
        # ax[0][0].set_title("VGG19, Cifar10")
        ax[0].set_xticks(range(0, 31, 5))
        ax[0].set_yticks(range(20, 91, 20))
        ax[0].set_ylabel("Accuracy (%)", fontsize=fontsize)
        ax[0].set_xlabel("Epoch", fontsize=fontsize)
        ax[0].tick_params(labelsize=fontsize)

        ax[1].plot(sa(self.baseline_loss, 71, 10)[::step], "-", label="Baseline")
        ax[1].plot(sa(self.quan2_loss, 71, 10)[::step], "--", label="2-bit quantization")
        ax[1].plot(sa(self.quan8_loss, 71, 10)[::step], ":", label="8-bit quantization")
        ax[1].plot(sa(self.sparse_loss, 71, 10)[::step], "-.", label="98% sparsification")
        ax[1].plot(sa(self.ssfdme_loss, 71, 10)[::step], ".", label="Ours")
        # ax[1].legend(prop={'size':6})
        ax[1].set_ylabel("Loss", fontsize=fontsize)
        ax[1].set_xlabel("Iteration (×$10^2$)", fontsize=fontsize)
        ax[1].tick_params(labelsize=fontsize)

        lines, labels = fig.axes[-1].get_legend_handles_labels()
        fig.legend(lines, labels, ncol=3, bbox_to_anchor=(0.8, 1.08), frameon=False, fontsize=fontsize)
        fig.savefig(f"./res/image/{self.model}_{self.dataset}.pdf", dpi=600, format='pdf', bbox_inches='tight')


if __name__ == '__main__':
    plot = Plot("VGG", "CIFAR")
